import TypeBuilder from "../baml_client/type_builder";
import { FieldType } from "../baml_client/type_builder";
import { b } from "./test-setup";

describe("Dynamic Type Tests", () => {
  describe("Import Tests", () => {
    it("should import TypeBuilder and FieldType from type_builder module", () => {
      // Test that TypeBuilder is importable
      const tb = new TypeBuilder();
      expect(tb).toBeInstanceOf(TypeBuilder);

      // Test that FieldType is importable and works
      expect(FieldType).toBeDefined();

      // Test that TypeBuilder methods return FieldType instances
      const stringType = tb.string();
      // TypeScript will ensure type safety at compile time
      // This test verifies runtime behavior
      expect(stringType).toBeDefined();
    });
  });

  describe("Basic Dynamic Types", () => {
    it("should work with dynamic types single", async () => {
      let tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().optional());
      tb.Person.addProperty("height", tb.float().optional()).description(
        "Height in meters",
      );
      tb.Hobby.addValue("CHESS");
      tb.Hobby.listValues().map(([name, v]) => v.alias(name.toLowerCase()));
      tb.Person.addProperty(
        "hobbies",
        tb.Hobby.type().list().optional(),
      ).description("Some suggested hobbies they might be good at");

      const res = await b.ExtractPeople(
        "My name is Harrison. My hair is black and I'm 6 feet tall. I'm pretty good around the hoop.",
        { tb },
      );
      expect(res.length).toBeGreaterThan(0);
    });

    it("should work with dynamic types enum", async () => {
      let tb = new TypeBuilder();
      const fieldEnum = tb.addEnum("Animal");
      const animals = ["giraffe", "elephant", "lion"];
      for (const animal of animals) {
        fieldEnum.addValue(animal.toUpperCase());
      }
      tb.Person.addProperty("animalLiked", fieldEnum.type());
      const res = await b.ExtractPeople(
        "My name is Harrison. My hair is black and I'm 6 feet tall. I'm pretty good around the hoop. I like giraffes.",
        { tb },
      );
      expect(res.length).toBeGreaterThan(0);
      expect(res[0]["animalLiked"]).toEqual("GIRAFFE");
    });

    it("should work with dynamic types class", async () => {
      let tb = new TypeBuilder();
      const animalClass = tb.addClass("Animal");
      animalClass
        .addProperty("animal", tb.string())
        .description("The animal mentioned, in singular form.");
      tb.Person.addProperty("animalLiked", animalClass.type());
      const res = await b.ExtractPeople(
        "My name is Harrison. My hair is black and I'm 6 feet tall. I'm pretty good around the hoop. I like giraffes.",
        { tb },
      );
      expect(res.length).toBeGreaterThan(0);
      const animalLiked = res[0]["animalLiked"];
      expect(animalLiked["animal"]).toContain("giraffe");
    });

    it("should work with dynamic literals", async () => {
      let tb = new TypeBuilder();
      const animals = tb.union(
        ["giraffe", "elephant", "lion"].map((animal) =>
          tb.literalString(animal.toUpperCase()),
        ),
      );
      tb.Person.addProperty("animalLiked", animals);
      const res = await b.ExtractPeople(
        "My name is Harrison. My hair is black and I'm 6 feet tall. I'm pretty good around the hoop. I like giraffes.",
        { tb },
      );
      expect(res.length).toBeGreaterThan(0);
      expect(res[0]["animalLiked"]).toEqual("GIRAFFE");
    });

    it("should work with dynamic inputs class", async () => {
      let tb = new TypeBuilder();
      tb.DynInputOutput.addProperty("new-key", tb.string().optional());

      const res = await b.DynamicInputOutput(
        { "new-key": "hi", testKey: "myTest" },
        { tb },
      );
      expect(res["new-key"]).toEqual("hi");
      expect(res["testKey"]).toEqual("myTest");
    });
  });

  describe("Complex Dynamic Types", () => {
    it("should work with dynamic inputs list", async () => {
      let tb = new TypeBuilder();
      tb.DynInputOutput.addProperty("new-key", tb.string().optional());

      const res = await b.DynamicListInputOutput(
        [{ "new-key": "hi", testKey: "myTest" }],
        { tb },
      );
      expect(res[0]["new-key"]).toEqual("hi");
      expect(res[0]["testKey"]).toEqual("myTest");
    });

    it("should work with dynamic output map", async () => {
      let tb = new TypeBuilder();
      tb.DynamicOutput.addProperty("hair_color", tb.string());
      tb.DynamicOutput.addProperty(
        "attributes",
        tb.map(tb.string(), tb.string()),
      ).description("Things like 'eye_color' or 'facial_hair'");

      const res = await b.MyFunc(
        "My name is Harrison. My hair is black and I'm 6 feet tall. I have blue eyes and a beard.",
        { tb },
      );

      expect(res.hair_color).toEqual("black");
      expect(res.attributes["eye_color"]).toEqual("blue");
      expect(res.attributes["facial_hair"]).toEqual("beard");
    });

    it("should work with dynamic output union", async () => {
      let tb = new TypeBuilder();

      const class1 = tb.addClass("Class1");
      class1.addProperty("meters", tb.float());

      const class2 = tb.addClass("Class2");
      class2.addProperty("feet", tb.float());
      class2.addProperty("inches", tb.float().optional());

      tb.DynamicOutput.addProperty(
        "height",
        tb.union([class1.type(), class2.type()]),
      );

      let res = await b.MyFunc(
        "My name is Harrison. My hair is black and I'm 6 feet tall.",
        { tb },
      );

      expect(res.height["feet"]).toEqual(6);

      res = await b.MyFunc(
        "My name is Harrison. My hair is black and I'm 1.8 meters tall.",
        { tb },
      );

      expect(res.height["meters"]).toEqual(1.8);
    });
  });

  describe("Add Baml", () => {
    it("should add to existing class", async () => {
      let tb = new TypeBuilder();
      tb.addBaml(`
        class ExtraPersonInfo {
            height int
            weight int
        }

        dynamic class Person {
            age int?
            extra ExtraPersonInfo?
        }
      `);
      let res = await b.ExtractPeople(
        "My name is John Doe. I'm 30 years old. I'm 6 feet tall and weigh 180 pounds. My hair is yellow.",
        { tb },
      );
      expect(res).toEqual([
        {
          name: "John Doe",
          age: 30,
          extra: { height: 6, weight: 180 },
          hair_color: "YELLOW",
        },
      ]);
    });

    it("should add to existing enum", async () => {
      let tb = new TypeBuilder();
      tb.addBaml(`
        dynamic enum Hobby {
          VideoGames
          BikeRiding
        }
      `);
      let res = await b.ExtractHobby("I play video games", { tb });
      expect(res).toEqual(["VideoGames"]);
    });

    it("should add both class and enum", async () => {
      let tb = new TypeBuilder();
      tb.addBaml(`
        class ExtraPersonInfo {
            height int @description("in feet")
            weight int @description("in pounds")
        }

        enum Job {
            Programmer
            Architect
            Musician
        }

        dynamic enum Hobby {
            VideoGames
            BikeRiding
        }

        dynamic enum Color {
            BROWN
        }

        dynamic class Person {
            age int?
            extra ExtraPersonInfo?
            job Job?
            hobbies Hobby[]
        }
      `);
      let res = await b.ExtractPeople(
        "My name is John Doe. I'm 30 years old. My height is 6 feet and I weigh 180 pounds. My hair is brown. I work as a programmer and enjoy bike riding.",
        { tb },
      );
      expect(res).toEqual([
        {
          name: "John Doe",
          age: 30,
          hair_color: "BROWN",
          job: "Programmer",
          hobbies: ["BikeRiding"],
          extra: { height: 6, weight: 180 },
        },
      ]);
    });

    it("should add baml with attrs", async () => {
      let tb = new TypeBuilder();
      tb.addBaml(`
        class ExtraPersonInfo {
            height int @description("In centimeters and rounded to the nearest whole number")
            weight int @description("In kilograms and rounded to the nearest whole number")
        }

        dynamic class Person {
            extra ExtraPersonInfo?
        }
      `);
      let res = await b.ExtractPeople(
        "My name is John Doe. I'm 30 years old. I'm 6 feet tall and weigh 180 pounds. My hair is yellow.",
        { tb },
      );
      expect(res).toEqual([
        {
          name: "John Doe",
          extra: { height: 183, weight: 82 },
          hair_color: "YELLOW",
        },
      ]);
    });
  });

  describe("TypeBuilder APIs", () => {
    it("should list properties", () => {
      let tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().list());
      tb.Person.addProperty("height", tb.float().optional()).description("Height in meters");

      const props = Object.fromEntries(tb.Person.listProperties());

      expect(props["last_name"].getType().equals(tb.string().list())).toBeTruthy();
      expect(props["height"].getType().equals(tb.float().optional())).toBeTruthy();
    });

    it("should reset", () => {
      let tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().list());
      tb.Person.addProperty("height", tb.float().optional()).description("Height in meters");
      tb.reset();

      const personPropsAfterTbReset = tb.Person.listProperties().map(([name, _]) => name);

      expect(personPropsAfterTbReset.includes("last_name")).toBeFalsy();
      expect(personPropsAfterTbReset.includes("height")).toBeFalsy();
    });

    it("should reset a class", () => {
      const tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().list());
      tb.Person.addProperty("height", tb.float().optional()).description("Height in meters");

      tb.DynamicOutput.addProperty("hair_color", tb.string());
      tb.DynamicOutput.addProperty("height", tb.float().optional()).description("Height in meters");

      tb.Person.reset();

      const personPropsAfterClassReset = tb.Person.listProperties().map(([name, _]) => name);
      const dynamicOutputPropsAfterClassReset = tb.DynamicOutput.listProperties().map(([name, _]) => name);

      expect(personPropsAfterClassReset.includes("last_name")).toBeFalsy();
      expect(personPropsAfterClassReset.includes("height")).toBeFalsy();

      expect(dynamicOutputPropsAfterClassReset.includes("hair_color")).toBeTruthy();
      expect(dynamicOutputPropsAfterClassReset.includes("height")).toBeTruthy();
    });

    it("should remove a property from a class", () => {
      const tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().list());
      tb.Person.addProperty("height", tb.float().optional()).description("Height in meters");

      tb.Person.removeProperty("last_name");

      const personProps = tb.Person.listProperties().map(([name, _]) => name);

      expect(personProps.includes("last_name")).toBeFalsy();
      expect(personProps.includes("height")).toBeTruthy();
    });

    it("should reset a dynamically added class", () => {
      const tb = new TypeBuilder();
      const personClass = tb.addClass("AddedPerson");
      personClass.addProperty("last_name", tb.string().list());
      personClass.addProperty("height", tb.float().optional()).description("Height in meters");

      personClass.reset();

      const personProps = personClass.listProperties().map(([name, _]) => name);

      expect(personProps.includes("last_name")).toBeFalsy();
      expect(personProps.includes("height")).toBeFalsy();
    });

    it("should remove a property from a dynamically added class", () => {
      const tb = new TypeBuilder();
      const personClass = tb.addClass("AddedPerson");
      personClass.addProperty("last_name", tb.string().list());
      personClass.addProperty("height", tb.float().optional()).description("Height in meters");

      personClass.removeProperty("last_name");
      const personProps = personClass.listProperties().map(([name, _]) => name);

      expect(personProps.includes("last_name")).toBeFalsy();
      expect(personProps.includes("height")).toBeTruthy();
    });

    it("should get property types from a class", () => {
      const tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().list());
      tb.Person.addProperty("height", tb.float().optional()).description("Height in meters");

      const props = Object.fromEntries(tb.Person.listProperties());

      expect(props["last_name"].getType().equals(tb.string().list())).toBeTruthy();
      expect(props["height"].getType().equals(tb.float().optional())).toBeTruthy();
    });

    it("should set property types", () => {
      const tb = new TypeBuilder();
      tb.Person.addProperty("last_name", tb.string().list());
      tb.Person.addProperty("height", tb.float().optional()).description("Height in meters");

      // Modify props
      let props = Object.fromEntries(tb.Person.listProperties());
      props["last_name"].setType(tb.string());
      props["height"].setType(tb.int());

      // Verify changes
      props = Object.fromEntries(tb.Person.listProperties());
      expect(props["last_name"].getType().equals(tb.string())).toBeTruthy();
      expect(props["height"].getType().equals(tb.int())).toBeTruthy();
    });
  });
});
